{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5YMscAJ4DCke"
   },
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/phoenix-logo-light.svg\" width=\"200\"/>\n",
    "        <br>\n",
    "        <a href=\"https://docs.arize.com/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://join.slack.com/t/arize-ai/shared_invite/zt-1px8dcmlf-fmThhDFD_V_48oU7ALan4Q\">Community</a>\n",
    "    </p>\n",
    "</center>\n",
    "<h1 align=\"center\">Tracing a LangChain and VertexAI Application</h1>\n",
    "\n",
    "LLM orchestration frameworks such as LangChain provide abstractions that enable users to build powerful applications in a few lines of code. However, the same abstractions can make it difficult to understand what is going on under the hood and to pinpoint the cause of issues.\n",
    "\n",
    "Phoenix makes your LLM applications *observable* by visualizing the underlying structure of each call to your chain and surfacing problematic \"spans\" of execution based on latency, token count, or other evaluation metrics.\n",
    "\n",
    "In this tutorial, you will:\n",
    "- Build a simple retrieval-augmented generation application over the Arize documentation using LangChain and VertexAI, in particular, using \"textembedding-gecko\" for embeddings and \"chat-bison\" for chat,\n",
    "- Record trace data in OpenInference format,\n",
    "- Inspect the traces and spans of your application to identify uncaught exceptions and sources of latency and cost.\n",
    "\n",
    "ℹ️️ This notebook requires access to the [Vertex AI API](https://cloud.google.com/vertex-ai/docs)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RTx42fVADCkh"
   },
   "source": [
    "## 1. Install Dependencies and Import Libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Downgrade Colab's pre-installed version of `shapely` for compatibility reasons.\n",
    "\n",
    "⚠️ If you run into a version compatibility error in a later cell, try restarting the runtime and re-running the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --force-reinstall \"shapely<2.0.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I4lNQeHLDCkh"
   },
   "source": [
    "Install Phoenix, LangChain, and the Google Cloud AI Platform SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"arize-phoenix[evals]\" google-api-python-client \"google-cloud-aiplatform[preview]\" langchain nest-asyncio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ht9MHETADCki"
   },
   "source": [
    "Import libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from urllib.request import urlopen\n",
    "\n",
    "import nest_asyncio\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import phoenix as px\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.chat_models import ChatVertexAI\n",
    "from langchain.embeddings import VertexAIEmbeddings\n",
    "from langchain.retrievers import KNNRetriever\n",
    "from phoenix.evals import (\n",
    "    HallucinationEvaluator,\n",
    "    QAEvaluator,\n",
    "    RelevanceEvaluator,\n",
    "    VertexAIModel,\n",
    "    run_evals,\n",
    ")\n",
    "from phoenix.session.evaluation import get_qa_with_reference, get_retrieved_documents\n",
    "from phoenix.trace import DocumentEvaluations, SpanEvaluations\n",
    "from phoenix.trace.langchain import LangChainInstrumentor\n",
    "from tqdm import tqdm\n",
    "\n",
    "nest_asyncio.apply()  # needed for concurrent evals in notebook environments\n",
    "\n",
    "try:\n",
    "    from google.colab.auth import authenticate_user\n",
    "\n",
    "    IS_COLAB = True\n",
    "except ImportError:\n",
    "    IS_COLAB = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6lt8OwKTDCkj"
   },
   "source": [
    "## 2. Set Configuration and Authenticate with Vertex AI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4l__rKB0K6JO"
   },
   "source": [
    "If you are running this notebook in Colab, sign in with your Gmail credentials. If running locally, you'll need to ensure that your `gcloud` is correctly configured to run Vertex AI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if IS_COLAB:\n",
    "    authenticate_user()\n",
    "else:\n",
    "    print(\n",
    "        \"If running locally, ensure that your gcloud is correctly configured to run with Vertex AI.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QkJWNxVxD0HS"
   },
   "source": [
    "Enter your project ID and location:\n",
    "* `project_id`: The default project to use when making API calls.\n",
    "* `location`: The default location to use when making API calls. If not set defaults to us-central-1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_id = input(\"Enter your GCP project ID and press enter:\\n\")\n",
    "location = input(\"Enter your GCP location and press enter:\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MhQJPCsyDCki"
   },
   "source": [
    "## 3. Launch Phoenix\n",
    "\n",
    "You can run Phoenix in the background to collect trace data emitted by any LlamaIndex application that has been instrumented with the `OpenInferenceTraceCallbackHandler`.\n",
    "\n",
    "Launch Phoenix and follow the instructions in the cell output to open the Phoenix UI (the UI should be empty because we have yet to run our LangChain application)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session = px.launch_app()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bxqNPYpQDCkj"
   },
   "source": [
    "## 4. Instantiate LangChain Instrumentor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LangChainInstrumentor().instrument()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S06RdfczDCkk"
   },
   "source": [
    "## 5. Build Your LLM Application\n",
    "\n",
    "Define a `RetrievalQA` chain leveraging \"textembedding-gecko\" and \"chat-bison\" from the VertexAI API. The knowledge base of this chain is built over the Arize documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = VertexAIEmbeddings(\n",
    "    model_name=\"textembedding-gecko\",\n",
    "    project=project_id,\n",
    "    location=location,\n",
    ")\n",
    "database_df = pd.read_parquet(\n",
    "    \"http://storage.googleapis.com/arize-phoenix-assets/datasets/unstructured/llm/context-retrieval/langchain-pinecone-vertexai/database.parquet\"\n",
    ")\n",
    "knn_retriever = KNNRetriever(\n",
    "    index=np.stack(database_df[\"text_vector\"]),\n",
    "    texts=database_df[\"text\"].tolist(),\n",
    "    embeddings=embeddings,\n",
    ")\n",
    "llm = ChatVertexAI(\n",
    "    model_name=\"chat-bison\",\n",
    "    project=project_id,\n",
    "    location=location,\n",
    ")\n",
    "chain = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    retriever=knn_retriever,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "s5_AWP3hDCkk"
   },
   "source": [
    "## 6. Run the Chain\n",
    "\n",
    "Download a small dataset of user queries to ask your application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"http://storage.googleapis.com/arize-phoenix-assets/datasets/unstructured/llm/context-retrieval/arize_docs_queries.jsonl\"\n",
    "queries = []\n",
    "with urlopen(url) as response:\n",
    "    for line in response:\n",
    "        line = line.decode(\"utf-8\").strip()\n",
    "        data = json.loads(line)\n",
    "        queries.append(data[\"query\"])\n",
    "queries[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4F6iOYePDCkk"
   },
   "source": [
    "Run your chain against ten queries. Notice that the tracer is attached in the `run` call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for query in tqdm(queries[:10]):\n",
    "    chain.invoke(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9IB25Bl9R02j"
   },
   "source": [
    "Check out your traces in Phoenix!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Open the Phoenix UI if you haven't already: {session.url}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Export and Evaluate Your Trace Data\n",
    "\n",
    "You can export your trace data as a pandas dataframe for further analysis and evaluation.\n",
    "\n",
    "In this case, we will export our `retriever` spans into two separate dataframes:\n",
    "- `queries_df`, in which the retrieved documents for each query are concatenated into a single column,\n",
    "- `retrieved_documents_df`, in which each retrieved document is \"exploded\" into its own row to enable the evaluation of each query-document pair in isolation.\n",
    "\n",
    "This will enable us to compute multiple kinds of evaluations, including:\n",
    "- relevance: Are the retrieved documents grounded in the response?\n",
    "- Q&A correctness: Are your application's responses grounded in the retrieved context?\n",
    "- hallucinations: Is your application making up false information?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "queries_df = get_qa_with_reference(px.Client())\n",
    "retrieved_documents_df = get_retrieved_documents(px.Client())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, define your evaluation model and your evaluators.\n",
    "\n",
    "Evaluators are built on top of language models and prompt the LLM to assess the quality of responses, the relevance of retrieved documents, etc., and provide a quality signal even in the absence of human-labeled data. Pick an evaluator type and instantiate it with the language model you want to use to perform evaluations using our battle-tested evaluation templates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_model = VertexAIModel(\n",
    "    project=project_id,\n",
    "    location=location,\n",
    "    model_name=\"text-bison\",\n",
    "    temperature=0.0,\n",
    ")\n",
    "hallucination_evaluator = HallucinationEvaluator(eval_model)\n",
    "qa_correctness_evaluator = QAEvaluator(eval_model)\n",
    "relevance_evaluator = RelevanceEvaluator(eval_model)\n",
    "\n",
    "hallucination_eval_df, qa_correctness_eval_df = run_evals(\n",
    "    dataframe=queries_df,\n",
    "    evaluators=[hallucination_evaluator, qa_correctness_evaluator],\n",
    "    provide_explanation=True,\n",
    ")\n",
    "relevance_eval_df = run_evals(\n",
    "    dataframe=retrieved_documents_df,\n",
    "    evaluators=[relevance_evaluator],\n",
    "    provide_explanation=True,\n",
    ")[0]\n",
    "\n",
    "px.Client().log_evaluations(\n",
    "    SpanEvaluations(eval_name=\"Hallucination\", dataframe=hallucination_eval_df),\n",
    "    SpanEvaluations(eval_name=\"QA Correctness\", dataframe=qa_correctness_eval_df),\n",
    "    DocumentEvaluations(eval_name=\"Relevance\", dataframe=relevance_eval_df),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your evaluations should now appear as annotations on the appropriate spans in Phoenix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"🚀 Open the Phoenix UI if you haven't already: {session.url}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Final Thoughts\n",
    "\n",
    "LLM Traces and the accompanying OpenInference Tracing specification is designed to be a category of telemetry data that is used to understand the execution of LLMs and the surrounding application context such as retrieval from vector stores and the usage of external tools such as search engines or APIs. It lets you understand the inner workings of the individual steps your application takes wile also giving you visibility into how your system is running and performing as a whole. \n",
    "\n",
    "LLM Evals are designed for simple, fast, and accurate LLM-based evaluations. They let you quickly benchmark the performance of your LLM application and help you identify the problematic spans of execution.\n",
    "\n",
    "For more details on Phoenix, LLM Tracing, and LLM Evals, checkout our [documentation](https://docs.arize.com/phoenix/)."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
